{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfa322bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "import cv2\n",
    "import os\n",
    "from PIL import Image\n",
    "from keras.preprocessing import image\n",
    "from keras.preprocessing.image import ImageDataGenerator\n",
    "\n",
    "def pre_processing(img):\n",
    "\n",
    "    return img / 127.5 - 1\n",
    "def get_data_gen_args(mode):\n",
    "    if mode == 'train' or mode == 'val':\n",
    "        x_data_gen_args = dict(preprocessing_function=pre_processing,\n",
    "                               shear_range=0.1,\n",
    "                               zoom_range=0.1,\n",
    "                               rotation_range=10,\n",
    "                               width_shift_range=0.1,\n",
    "                               height_shift_range=0.1,\n",
    "                               fill_mode='constant',\n",
    "                               horizontal_flip=True)\n",
    "        y_data_gen_args = dict(shear_range=0.1,\n",
    "                               zoom_range=0.1,\n",
    "                               rotation_range=10,\n",
    "                               width_shift_range=0.1,\n",
    "                               height_shift_range=0.1,\n",
    "                               fill_mode='constant',\n",
    "                               horizontal_flip=True)\n",
    "\n",
    "\n",
    "    elif mode == 'test':\n",
    "        x_data_gen_args = dict(preprocessing_function=pre_processing)\n",
    "        y_data_gen_args = dict()\n",
    "    else:\n",
    "        print(\"Data_generator function should get mode arg 'train' or 'val' or 'test'.\")\n",
    "        return -1\n",
    "\n",
    "    return x_data_gen_args, y_data_gen_args\n",
    "\n",
    "\n",
    "def get_data_gen_args_cls(mode):\n",
    "    if mode == 'train' or mode == 'val':\n",
    "        x_data_gen_args = dict(preprocessing_function=pre_processing,\n",
    "                               shear_range=0.1,\n",
    "                               zoom_range=0.1,\n",
    "                               rotation_range=10,\n",
    "                               width_shift_range=0.1,\n",
    "                               height_shift_range=0.1,\n",
    "                               fill_mode='constant',\n",
    "                               horizontal_flip=True)\n",
    "\n",
    "    elif mode == 'test':\n",
    "        x_data_gen_args = dict(preprocessing_function=pre_processing)\n",
    "    else:\n",
    "        print(\"Data_generator function should get mode arg 'train' or 'val' or 'test'.\")\n",
    "        return -1\n",
    "\n",
    "    return x_data_gen_args\n",
    "\n",
    "def get_data_gen_args_seg(mode):\n",
    "    if mode == 'train' or mode == 'val':\n",
    "        y_data_gen_args = dict(shear_range=0.1,\n",
    "                               zoom_range=0.1,\n",
    "                               rotation_range=10,\n",
    "                               width_shift_range=0.1,\n",
    "                               height_shift_range=0.1,\n",
    "                               fill_mode='constant',\n",
    "                               horizontal_flip=True)\n",
    "\n",
    "    elif mode == 'test':\n",
    "        y_data_gen_args = dict()\n",
    "    else:\n",
    "        print(\"Data_generator function should get mode arg 'train' or 'val' or 'test'.\")\n",
    "        return -1\n",
    "\n",
    "    return y_data_gen_args\n",
    "\n",
    "def data_generator_pair(nameList, segLabel, cLabel, image_shape, nb_class, b_size, mode):\n",
    "    # Make ImageDataGenerator.\n",
    "    x_data_gen_args, y_data_gen_args = get_data_gen_args(mode)\n",
    "    x_data_gen = ImageDataGenerator(**x_data_gen_args)\n",
    "    y_data_gen = ImageDataGenerator(**y_data_gen_args)\n",
    "    First = True\n",
    "    d_size = len(nameList)\n",
    "    shuffled_idx = list(range(d_size))\n",
    "    y_seg = []\n",
    "    while True:\n",
    "        random.shuffle(shuffled_idx)\n",
    "        for i in range(d_size):\n",
    "            idx = shuffled_idx[i]\n",
    "            name = nameList[idx]\n",
    "            L = int(cLabel[idx])\n",
    "            img = Image.open(name)\n",
    "            y_img = load_data(segLabel[idx], image_shape, mode=\"label\")\n",
    "            img_rgb = img.convert('RGB')\n",
    "            img_rgb = img_rgb.resize((128, 128), Image.ANTIALIAS)\n",
    "            img_rgb_arr = np.array(img_rgb)\n",
    "            img_rgb_arr = np.reshape(img_rgb_arr, [1, 128, 128, 3])\n",
    "            y = np.zeros([1, nb_class])\n",
    "            y[0, L] = 1\n",
    "            if First:\n",
    "                X_train = img_rgb_arr\n",
    "                y_train_one = y\n",
    "                First = False\n",
    "            else:\n",
    "                X_train = np.concatenate((X_train, img_rgb_arr), axis=0)\n",
    "                y_train_one = np.concatenate((y_train_one, y), axis=0)\n",
    "            y_seg.append(y_img)\n",
    "            if X_train.shape[0] == b_size:\n",
    "                y_try = np.argmax(y_train_one, axis=1)\n",
    "                seed = random.randrange(1, 1000)\n",
    "\n",
    "                x_tmp_gen = x_data_gen.flow(X_train, y_try,\n",
    "                                            batch_size=b_size,\n",
    "                                            seed=seed)\n",
    "                y_tmp_gen = y_data_gen.flow(np.array(y_seg), y_try,\n",
    "                                            batch_size=b_size,\n",
    "                                            seed=seed)\n",
    "                x_result, y_try1 = next(x_tmp_gen)\n",
    "                y_result, y_try2 = next(y_tmp_gen)\n",
    "                y_train_one1 = np.zeros([y_train_one.shape[0], y_train_one.shape[1]])\n",
    "                for j in range(b_size):\n",
    "                    y_train_one1[j, y_try1[j]] = 1.0\n",
    "\n",
    "                # augmented\n",
    "                yield [x_result, x_result], [y_train_one1, binarylab(b_size, y_result, image_shape, nb_class)]\n",
    "                # original\n",
    "                # yield X_train, y_train_one\n",
    "                First = True\n",
    "                y_seg.clear()\n",
    "\n",
    "def data_generator_pair_different_dataset(cls_nameList, seg_nameList, seg_gt, cLabel, image_shape, nb_class, b_size, mode):\n",
    "    # Make ImageDataGenerator.\n",
    "\n",
    "    cls_train_generator = data_generator_cls(nameList=cls_nameList, cLabel=cLabel, image_shape=image_shape, nb_class=nb_class, b_size=1, mode=mode)\n",
    "    seg_data_gen = data_generator_seg(nameList=seg_nameList, segLabel=seg_gt, img_shape=image_shape, nb_class=nb_class, b_size=1, mode=mode)\n",
    "    while True:\n",
    "        cls_imgs = []\n",
    "        cls_labels = []\n",
    "        for k in range(b_size):\n",
    "            cls_data = cls_train_generator.__next__()\n",
    "            cls_imgs.append(cls_data[0][0])\n",
    "            cls_labels.append(cls_data[1][0])\n",
    "        cls_imgs = np.array(cls_imgs)\n",
    "        cls_labels = np.array(cls_labels)\n",
    "        # print('cls_img shape:{}'.format(cls_imgs.shape))\n",
    "        # print('cls_img label shape:{}'.format(cls_labels.shape))\n",
    "        seg_datas = []\n",
    "        seg_labels = []\n",
    "        for k in range(b_size):\n",
    "            seg_data_lb = seg_data_gen.__next__()\n",
    "            seg_datas.append(seg_data_lb[0][0])\n",
    "            seg_labels.append(seg_data_lb[1][0])\n",
    "        seg_datas = np.array(seg_datas)\n",
    "        seg_labels = np.array(seg_labels)\n",
    "        yield [cls_imgs, seg_datas], [cls_labels, seg_labels]\n",
    "\n",
    "def data_generator_cls(nameList, cLabel, image_shape, nb_class, b_size, mode):\n",
    "    # Make ImageDataGenerator.\n",
    "    x_data_gen_args = get_data_gen_args_cls(mode)\n",
    "    x_data_gen = ImageDataGenerator(**x_data_gen_args)\n",
    "    First = True\n",
    "    d_size = len(nameList)\n",
    "    shuffled_idx = list(range(d_size))\n",
    "    y_seg = []\n",
    "    while True:\n",
    "        random.shuffle(shuffled_idx)\n",
    "        for i in range(d_size):\n",
    "            idx = shuffled_idx[i]\n",
    "            name = nameList[idx]\n",
    "            L = int(cLabel[idx])\n",
    "            img = Image.open(name)\n",
    "            img_rgb = img.convert('RGB')\n",
    "            img_rgb = img_rgb.resize((image_shape[1], image_shape[0]), Image.ANTIALIAS)\n",
    "            img_rgb_arr = np.array(img_rgb)\n",
    "            img_rgb_arr = np.reshape(img_rgb_arr, [1, image_shape[1], image_shape[0], 3])\n",
    "            y = np.zeros([1, nb_class])\n",
    "            y[0, L] = 1\n",
    "            if First:\n",
    "                X_train = img_rgb_arr\n",
    "                y_train_one = y\n",
    "                First = False\n",
    "            else:\n",
    "                X_train = np.concatenate((X_train, img_rgb_arr), axis=0)\n",
    "                y_train_one = np.concatenate((y_train_one, y), axis=0)\n",
    "            if X_train.shape[0] == b_size:\n",
    "                y_try = np.argmax(y_train_one, axis=1)\n",
    "                seed = random.randrange(1, 1000)\n",
    "\n",
    "                x_tmp_gen = x_data_gen.flow(X_train, y_try,\n",
    "                                            batch_size=b_size,\n",
    "                                            seed=seed)\n",
    "\n",
    "                x_result, y_try1 = next(x_tmp_gen)\n",
    "                y_train_one1 = np.zeros([y_train_one.shape[0], y_train_one.shape[1]])\n",
    "                for j in range(b_size):\n",
    "                    y_train_one1[j, y_try1[j]] = 1.0\n",
    "\n",
    "                # augmented\n",
    "                yield x_result, y_train_one1\n",
    "                # original\n",
    "                # yield X_train, y_train_one\n",
    "                First = True\n",
    "                y_seg.clear()\n",
    "\n",
    "def data_generator_seg(nameList, segLabel, img_shape, nb_class, b_size, mode):\n",
    "\n",
    "    # Make ImageDataGenerator.\n",
    "    x_data_gen_args, y_data_gen_args = get_data_gen_args(mode)\n",
    "    x_data_gen = ImageDataGenerator(**x_data_gen_args)\n",
    "    y_data_gen = ImageDataGenerator(**y_data_gen_args)\n",
    "\n",
    "    # random index for random data access.\n",
    "    d_size = len(nameList)\n",
    "    shuffled_idx = list(range(d_size))\n",
    "\n",
    "    x = []\n",
    "    y = []\n",
    "    while True:\n",
    "        random.shuffle(shuffled_idx)\n",
    "        for i in range(d_size):\n",
    "            idx = shuffled_idx[i]\n",
    "            x_img = load_data(nameList[idx], img_shape, mode=\"data\")\n",
    "            y_img = load_data(segLabel[idx], img_shape, mode=\"label\")\n",
    "\n",
    "            x.append(x_img)\n",
    "            y.append(y_img)\n",
    "\n",
    "            if len(x) == b_size:\n",
    "                # Adapt ImageDataGenerator flow method for data augmentation.\n",
    "                _ = np.zeros(b_size)\n",
    "                seed = random.randrange(1, 1000)\n",
    "\n",
    "                x_tmp_gen = x_data_gen.flow(np.array(x), _,\n",
    "                                            batch_size=b_size,\n",
    "                                            seed=seed)\n",
    "                y_tmp_gen = y_data_gen.flow(np.array(y), _,\n",
    "                                            batch_size=b_size,\n",
    "                                            seed=seed)\n",
    "\n",
    "                # Finally, yield x, y data.\n",
    "                x_result, _ = next(x_tmp_gen)\n",
    "                y_result, _ = next(y_tmp_gen)\n",
    "\n",
    "                yield x_result, binarylab(b_size, y_result, img_shape, nb_class)\n",
    "\n",
    "                x.clear()\n",
    "                y.clear()\n",
    "\n",
    "\n",
    "def load_data(path, img_shape, mode=None):\n",
    "    img = Image.open(path)\n",
    "    img = img.resize((img_shape[1],img_shape[0]))\n",
    "    #w,h = img.size\n",
    "    #if w < h:\n",
    "    #    if w < size:\n",
    "    #        img = img.resize((size, size*h//w))\n",
    "    #        w, h = img.size\n",
    "    #else:\n",
    "    #    if h < size:\n",
    "    #        img = img.resize((size*w//h, size))\n",
    "    #        w, h = img.size\n",
    "    #img = img.crop((int((w-size)*0.5), int((h-size)*0.5), int((w+size)*0.5), int((h+size)*0.5)))\n",
    "    #img.show()\n",
    "    if mode==\"original\":\n",
    "        return img\n",
    "\n",
    "    if mode==\"label\":\n",
    "        y = np.array(img, dtype=np.int32)\n",
    "        mask = y == 255\n",
    "        y[mask] = 21\n",
    "        #y = binarylab(y, size, 21)\n",
    "        y = np.expand_dims(y, axis=-1)\n",
    "        return y\n",
    "    if mode==\"data\":\n",
    "        img = img.convert('RGB')\n",
    "        X = image.img_to_array(img)\n",
    "        #X = np.expand_dims(X, axis=0)\n",
    "        #X = preprocess_input(X)\n",
    "        return X\n",
    "\n",
    "def binarylab(b_size, y_img, img_shape, nb_class):\n",
    "    y_img = np.squeeze(y_img, axis=3)\n",
    "    result_map = np.zeros((b_size, img_shape[0], img_shape[1], nb_class))\n",
    "\n",
    "    # For np.where calculation.\n",
    "    for i in range(nb_class):\n",
    "        mask = (y_img == i)\n",
    "        result_map[:, :, :, i] = np.where(mask, 1, 0)\n",
    "\n",
    "    return result_map\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
